# coding=utf-8

import pyshark
import dpkt
from dpkt.ip import IP
from dpkt.tcp import TCP
from dpkt.udp import UDP
#from util.LinuxCookedProcess import linux_cooked_process
#from IPStream import IPStream
from TransportStream import TransportStream
import sys
import time
import socket
import os
import re
import csv
import math
import numpy as np

#client_ip = '192.168.12.125'     # 这里采用简单的方法确定客户端，客户端地址为私有地址
#client_ip = '10.20'        # 这里采用简单的方法确定客户端，客户端地址为私有地址

def calculation(arr):
    mean_,min_,max_,std_=0,0,0,0
    if len(arr) < 1:
        return mean_,min_,max_,std_
    elif len(arr) == 1:
        return arr[0],arr[0],arr[0],0
    else:
        min_=round(min(arr),14)
        max_=round(max(arr),14)
        mean_ = round(sum(arr)/len(arr),14)
        sd = sum([(i - mean_) ** 2 for i in arr])
        std_ = round(math.sqrt(sd / (len(arr)-1)),14)
        
    return mean_,min_,max_,std_
    
def packet_iat(arr):
    piat=[]
    if len(arr)>0:
        pre_time = arr[0]
        for pkt in arr[1:]:
            next_time = pkt
            piat.append(next_time-pre_time)
            pre_time=next_time
        piat_mean,piat_min,piat_max,piat_std=calculation(piat)
    else:
        piat_mean,piat_min,piat_max,piat_std=0,0,0,0
    return piat_mean,piat_min,piat_max,piat_std
    
# 根据规则区分服务器和客户端
def NormalizationSrcDst(src,sport,dst,dport):
    if sport < dport:
        return (dst,dport,src,sport)
    elif sport == dport:
        src_ip = "".join(src.split('.'))
        dst_ip = "".join(dst.split('.'))
        if int(src_ip) < int(dst_ip):
            return (dst,dport,src,sport)
        else:
            return (src,sport,dst,dport)
    else:
        return (src,sport,dst,dport)


class Extract_Flow_Features(object):
    """
    借助于dkpt，从数据包里提取出流特征
    """
    def __init__(self,pcap_path):
        """
        extract_flow_features对象构造函数
        :param pcap_path:
        """
        self.pcap_path = pcap_path

    def extract(self):
        #ip_packet_dict = dict()
        trans_packet_dict = dict()
        try:
            # print("Read pcap started." + str(time.time()))
            # print(self.pcap_path)
            start = None
            deal_pacp_path = self.pcap_path[:]
   
            # if linux_cooked_process(deal_pacp_path):
            #     deal_pacp_path += ".deal.pcap"
                # print("linux")
            with open(deal_pacp_path, "rb") as deal_pcap:
                # print("open")
                pcap = dpkt.pcap.Reader(deal_pcap)
                for ts, buf in pcap:
                    #print("ts:",ts)
                    if not start:
                        start = ts
                    try:
                        ether = dpkt.ethernet.Ethernet(buf)
                        ip = ether.data
                       # if ether.type == dpkt.ethernet.ETH_TYPE_IP:
                       #     print(type(ip))
                    except Exception as e:
                        #print(e,self.pcap_path[:])
                        continue
                    ts = ts - start
                    #print("ts:", ts)
                    #print(ip)
                    if not isinstance(ip, IP):
                        continue
                    if ip:
                        '''
                        if client_ip in socket.inet_ntoa(ip.src):
                            src, dst = socket.inet_ntoa(ip.src), socket.inet_ntoa(ip.dst)
                        else:
                            src, dst = socket.inet_ntoa(ip.dst), socket.inet_ntoa(ip.src)
                        
                        '''
                        ip_payload = ip.data  # 判断传输层协议
                        if isinstance(ip_payload, TCP):
                            trans_protocol = 6  # 判断为tcp
                        elif isinstance(ip_payload, UDP):
                            trans_protocol = 17  # 判断为udp
                        else:
                            continue
                        trans = ip.data
                        src,sport,dst,dport = NormalizationSrcDst(socket.inet_ntoa(ip.src),trans.sport,socket.inet_ntoa(ip.dst),trans.dport)

                        '''
                        if (src, dst) not in ip_packet_dict.keys():
                            ip_stream = IPStream(src, dst, ts)
                            ip_packet_dict[(src, dst)] = ip_stream
                        else:
                            ip_stream = ip_packet_dict[(src, dst)]
                            if ts > ip_stream.end_time:
                                ip_stream.end_time = ts
                            elif ts < ip_stream.start_time:
                                ip_stream.start_time = ts
                        ip_stream.pkts += 1  # 累计ip总包数
                        ip_stream.bytes += len(buf)  # 累计ip总字节数
                        #if client_ip in socket.inet_ntoa(ip.src):
                        if src == socket.inet_ntoa(ip.src):
                            ip_stream.fwd_pkts += 1  # 累计ip正向包数
                            ip_stream.fwd_bytes += len(buf)  # 累计ip正向字节数
                        else:
                            ip_stream.bwd_pkts += 1  # 累计ip反向包数
                            ip_stream.bwd_bytes += len(buf)  # 累计ip反向字节数
                        ip_stream.ip_datagrams.append((ts, ip))
                        '''
                        
                        '''
                        ip_payload = ip.data  # 判断传输层协议
                        if isinstance(ip_payload, TCP):
                            trans_protocol = 6  # 判断为tcp
                        elif isinstance(ip_payload, UDP):
                            trans_protocol = 17  # 判断为udp
                        else:
                            continue
                        
                        trans = ip.data  # 传输层数据
                        if client_ip in socket.inet_ntoa(ip.src):
                            sport, dport = trans.sport, trans.dport
                        else:
                            sport, dport = trans.dport, trans.sport
                        '''
                        if (src, dst, sport, dport, trans_protocol) not in trans_packet_dict.keys():
                            #trans_stream = TransportStream(src, dst, sport, dport, trans_protocol, ip_stream, ts)
                            trans_stream = TransportStream(src, dst, sport, dport, trans_protocol,  ts)
                            trans_packet_dict[(src, dst, sport, dport, trans_protocol)] = trans_stream
                            '''
                            if trans_protocol == 6:
                                ip_stream.tcp_streams.append(trans_stream)
                            else:
                                ip_stream.udp_streams.append(trans_stream)
                            '''
                        else:
                            trans_stream = trans_packet_dict[(src, dst, sport, dport, trans_protocol)]
                            if ts > trans_stream.end_time:
                                trans_stream.end_time = ts
                            elif ts < trans_stream.start_time:
                                trans_stream.start_time = ts
                        #print(len(trans_packet_dict))
                        trans_stream.pkts += 1
                        trans_stream.bytes += len(trans.data)
                        trans_stream.pkt_size.append(len(trans.data))
                        trans_stream.pkt_time.append(ts)
                        
                        #if client_ip in socket.inet_ntoa(ip.src):
                        if src == socket.inet_ntoa(ip.src):
                            trans_stream.fwd_pkts += 1
                            trans_stream.fwd_bytes += len(trans.data)
                            trans_stream.fwd_pkt_size.append(len(trans.data))
                            trans_stream.fwd_pkt_time.append(ts)
                        else:
                            trans_stream.bwd_pkts += 1
                            trans_stream.bwd_bytes += len(trans.data)
                            trans_stream.bwd_pkt_size.append(len(trans.data))
                            trans_stream.bwd_pkt_time.append(ts)
                        if trans_protocol == 6:
                            trans_stream.tcp_packets.append((ts, trans))
                        else:
                            trans_stream.udp_packets.append((ts, trans))

            for flow in trans_packet_dict.values():
                #print(flow)
                pkt_size_mean, pkt_size_min, pkt_size_max, pkt_size_std = calculation(flow.pkt_size)
                fwd_pkt_size_mean, fwd_pkt_size_min, fwd_pkt_size_max, fwd_pkt_size_std = calculation(flow.fwd_pkt_size)
                bwd_pkt_size_mean, bwd_pkt_size_min, bwd_pkt_size_max, bwd_pkt_size_std = calculation(flow.bwd_pkt_size)
                pkt_iat_mean, pkt_iat_min, pkt_iat_max, pkt_iat_std = packet_iat(flow.pkt_time)
                fwd_iat_mean, fwd_iat_min, fwd_iat_max, fwd_iat_std = packet_iat(flow.fwd_pkt_time)
                bwd_iat_mean, bwd_iat_min, bwd_iat_max, bwd_iat_std = packet_iat(flow.bwd_pkt_time)
               
                feature = [flow.src, flow.dst, flow.sport, flow.dport, flow.trans_protocol, flow.duration, flow.pkts, flow.fwd_pkts, flow.bwd_pkts, flow.bytes, flow.fwd_bytes, flow.bwd_bytes,pkt_size_mean, pkt_size_min, pkt_size_max, pkt_size_std, fwd_pkt_size_mean, fwd_pkt_size_min, fwd_pkt_size_max, fwd_pkt_size_std, bwd_pkt_size_mean, bwd_pkt_size_min, bwd_pkt_size_max, bwd_pkt_size_std, flow.flow_rate, flow.pkt_number_mean, flow.fwd_flow_rate, flow.fwd_pkt_number_mean, flow.bwd_flow_rate, flow.bwd_pkt_number_mean, pkt_iat_mean, pkt_iat_min, pkt_iat_max, pkt_iat_std, fwd_iat_mean, fwd_iat_min, fwd_iat_max, fwd_iat_std, bwd_iat_mean, bwd_iat_min, bwd_iat_max, bwd_iat_std]
                #print(feature)
                file = open(csvname,"a+")
                writer = csv.writer(file)
                writer.writerow(feature)
                file.close()

            
                    
            # print("共有" + str(len(trans_packet_dict.keys()))+"条流，其中：")
            # print("tcp流的个数为"+str(tcp_len)+' '+"udp流的个数为"+str(udp_len))
            # print("业务流信息为 "+"src:"+str(largest_bytes_flow.src)+" dst:"+
            #       str(largest_bytes_flow.dst)+" sport:"+str(largest_bytes_flow.sport)+
            #       " dport:"+str(largest_bytes_flow.dport)+" protocol:"+str(largest_bytes_flow.trans_protocol)+
            #       " pkts:"+str(largest_bytes_flow.pkts)+" bytes:"+str(largest_bytes_flow.bytes)+
            #       " pkt_len_mean:"+str(largest_bytes_flow.pkt_len_mean))
            # print("组流成功")
            #return largest_bytes_flow
           
        except Exception as e:
            print(e,self.pcap_path[:])
            return [], []

feature_name = ['Src IP','Dst IP','Src Port','Dst Port','Protocol','Flow Duration','Packet Number','Fwd Packet Number','Bwd Packet Number','Total Bytes','Total Fwd Bytes','Total Bwd Bytes','Packet Length Mean','Packet Length Min','Packet Length Max','Packet Length Std','Fwd Packet Length Mean','Fwd Packet Length Min','Fwd Packet Length Max','Fwd Packet Length Std','Bwd Packet Length Mean','Bwd Packet Length Min','Bwd Packet Length Max','Bwd Packet Length Std','Flow Byte/s','Flow Packet/s','Fwd Flow Byte/s','Fwd Flow Packet/s','Bwd Flow Byte/s','Bwd Flow Packet/s','Packet Iat Mean','Packet Iat Min','Packet Iat Max','Packet Iat Std','Fwd Iat Mean','Fwd Iat Min','Fwd Iat Max','Fwd Iat Std','Bwd Iat Mean','Bwd Iat Min','Bwd Iat Max','Bwd Iat Std']

csvname = "lol.csv"

if __name__ == "__main__":
    file_dir = '/Users/quanjingwen/Desktop/pcap'           # pcap包的文件夹
    with open(csvname,"w+") as file:
        writer = csv.writer(file)
        print("write colname")
        writer.writerow(feature_name)
    for root, dirs, files in os.walk(file_dir):         # 遍历每个pcap文件
        for pcap_path in files:
            pcap_path = os.path.join(root, pcap_path)  # 得到pcap文件的地址
            print(pcap_path)
            Extract_Flow_Features(pcap_path).extract()





